import numpy as np
import torch
from torch import nn
from d2l import torch as d2l
import pandas as pd

#import库 以及 用到的函数定义见4.10.2
k, num_epochs, weight_decay, batch_size = 5, 100, 0, 64
animator_lr = d2l.Animator(xlabel='lr', ylabel='rmse', xlim=[0,20], ylim=[0.0, 0.2],
                        legend=[ 'train ', 'valid ','mean(train+valid)'])
for lr in np.linspace(0.,20,20):
    train_l, valid_l = k_fold(loss,get_net(),k, train_features, train_labels, num_epochs, lr,
                          weight_decay, batch_size,show=0)
    print(f'lr={float(lr):f}  {k}-折验证: 平均训练log rmse: {float(train_l):f}, '
            f'平均验证log rmse: {float(valid_l):f}')
    animator_lr.add(lr, (train_l, valid_l,0.5*(train_l+valid_l)))